import torch.nn as nn
import torch
from torchsummary import summary

class DownSample(nn.Module):
    def __init__(self, in_channels, out_channels, stride=2):
        super(DownSample, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1), stride=stride, padding=0)
        self.batch_normal = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        x = self.batch_normal(x)
        return x

class SELayer(nn.Module):
    def __init__(self, channel, reduction=64):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)
class Res2NetBottleNeck(nn.Module):
    def __init__(self, in_channels, latent_channels, out_channels, stride=1, scale=4, group=1):
        super(Res2NetBottleNeck, self).__init__()
        if latent_channels % group != 0:
            raise ValueError('group must be divisible by latent_channels')
        self.scale = scale
        self.conv1 = nn.Conv2d(in_channels, latent_channels, kernel_size=(1, 1), stride=stride, padding=0)
        self.bn1 = nn.BatchNorm2d(latent_channels)
        self.conv2 = nn.ModuleList(
            [nn.Conv2d(latent_channels // scale, latent_channels // scale, kernel_size=(3, 3), stride=(1, 1), padding=1,
                       groups=group) for _ in range(scale - 1)])
        self.bn2 = nn.ModuleList([nn.BatchNorm2d(latent_channels // scale) for _ in range(scale - 1)])
        self.relu = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(latent_channels, out_channels, kernel_size=(1, 1), stride=(1, 1), padding=0)
        self.se_layer = SELayer(out_channels)
        if in_channels != out_channels:
            self.down_sample = DownSample(in_channels, out_channels, stride=stride)
        else:
            self.down_sample = None



    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        xs = torch.chunk(out, self.scale, dim=1)
        ys = []
        for s in range(self.scale):
            if s == 0:
                ys.append(xs[s])
            elif s == 1:
                ys.append(self.relu(self.bn2[s - 1](self.conv2[s - 1](xs[s]))))
            else:
                ys.append(self.relu(self.bn2[s - 1](self.conv2[s - 1](xs[s] + ys[-1]))))
        out = torch.cat(ys, dim=1)
        out = self.conv3(out)
        out = self.se_layer(out)
        if self.down_sample is not None:
            x = self.down_sample(x)
        out = self.relu(out + x)
        return out


class Res2Net(nn.Module):
    def __init__(self):
        super(Res2Net, self).__init__()
        # format: H×W×C
        # 32×32×3 -> 32×32×32
        self.conv = nn.Conv2d(3, 32, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3))
        self.bn = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        # 32×32×32 -> 32×32×32
        # parm: in_channels, out_channels, stride, down_sample=False)
        # self.basic1 = BasicBlock(32, 32, 1)
        # parm: in_channels, latent_channels, out_channels, scale=4, group=1)
        self.bottle_neck1 = Res2NetBottleNeck(32, 64, 32)
        # 32×32×32 -> 16×16×64
        # self.basic2 = BasicBlock(32, 64, 2, True)
        self.bottle_neck2 = Res2NetBottleNeck(32, 128, 64, 2)
        # 16×16×64 -> 16×16×64
        # self.basic3 = BasicBlock(64, 64, 1)
        self.bottle_neck3 = Res2NetBottleNeck(64, 128, 64)
        # 16×16×64 -> 8×8×128
        # self.basic4 = BasicBlock(64, 128, 2, True)
        self.bottle_neck4 = Res2NetBottleNeck(64, 256, 128, 2)
        # 8×8×128 -> 8×8×128
        self.fc1 = nn.Linear(8192, 1024)
        self.fc2 = nn.Linear(1024, 100)
        self.relu2 = nn.ReLU()

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu1(x)
        x = self.bottle_neck1(x)
        x = self.bottle_neck2(x)
        x = self.bottle_neck3(x)
        x = self.bottle_neck4(x)
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        x = self.relu2(x)
        return self.fc2(x)

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    net = Res2Net().to(device)
    input_size = (3, 32, 32)
    summary(net, input_size=input_size)
    x = torch.randn(2, 3, 32, 32).to(device)
    y = net(x)
    print(y.shape)